import os
from typing import Any, Dict, List, Tuple

import faiss
import numpy as np
from subway_qa.adapters import embedding


def build_index(
    qa_pairs: List[Dict[str, Any]],
    embeddings: np.ndarray,
    index_path: str = "qa_index.faiss",
) -> None:
    """
    Index QA pairs and their embeddings using FAISS.

    Args:
        qa_pairs: List of dictionaries containing QA pairs
        embeddings: Numpy array of embedding vectors
        index_path: Path to save the FAISS index
    """
    # Get dimension from the embeddings
    dimension = embeddings.shape[1]

    # Create a FAISS index
    index = faiss.IndexFlatL2(dimension)

    # Add the embeddings to the index
    index.add(embeddings.astype(np.float32))

    # Save the index
    faiss.write_index(index, index_path)

    # Save the QA pairs separately
    np.save(index_path + ".qa_pairs", np.array(qa_pairs, dtype=object))


if os.path.exists("qa_index.faiss"):
    index = faiss.read_index("qa_index.faiss")
else:
    index = None

if os.path.exists("qa_index.faiss.qa_pairs.npy"):
    qa_pairs = np.load("qa_index.faiss.qa_pairs.npy", allow_pickle=True).tolist()
else:
    qa_pairs = None


def search(query: str, k: int = 5) -> Tuple[List[Dict[str, Any]], List[float]]:
    """
    Search for similar QA pairs using a query embedding.

    Args:
        query: query text
        k: Number of results to return

    Returns:
        Tuple of (list of QA pairs, list of distances)
    """
    query_embedding = np.array(embedding(query))

    # Ensure the query embedding is a 2D array
    if len(query_embedding.shape) == 1:
        query_embedding = query_embedding.reshape(1, -1)

    # Perform the search
    distances, indices = index.search(query_embedding.astype(np.float32), k)

    # Get the corresponding QA pairs
    result_qa_pairs = [qa_pairs[idx] for idx in indices[0]]

    return result_qa_pairs


if __name__ == "__main__":
    print(search("地铁票价是怎么计算的？"))